#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
3D彩票区域转换特征工具类
为所有算法提供统一的区域转换特征计算功能
"""

import numpy as np
import pandas as pd
import torch

class Data3DZoneFeatureExtractor:
    """3D彩票区域转换特征提取器"""
    
    def __init__(self):
        """初始化特征提取器"""
        pass
    
    def extract_zone_transitions(self, data_sequence):
        """
        从数据序列中提取区域转换特征
        
        Args:
            data_sequence: 数据序列，形状为 [window_size, 3] 或 [batch_size, window_size, 3]
            
        Returns:
            包含区域转换特征的数据，形状为 [window_size, 6] 或 [batch_size, window_size, 6]
        """
        # 确保输入是numpy数组
        if isinstance(data_sequence, torch.Tensor):
            data_sequence = data_sequence.cpu().numpy()
        
        # 处理不同的输入形状
        if len(data_sequence.shape) == 2:
            # 单个序列 [window_size, 3]
            return self._extract_single_sequence_zone_features(data_sequence)
        elif len(data_sequence.shape) == 3:
            # 批量序列 [batch_size, window_size, 3]
            batch_features = []
            for seq in data_sequence:
                batch_features.append(self._extract_single_sequence_zone_features(seq))
            return np.array(batch_features)
        else:
            raise ValueError(f"不支持的数据形状: {data_sequence.shape}")
    
    def _extract_single_sequence_zone_features(self, sequence):
        """
        从单个序列中提取区域转换特征
        
        Args:
            sequence: 单个序列，形状为 [window_size, 3]
            
        Returns:
            包含区域转换特征的序列，形状为 [window_size, 6]
        """
        window_size, num_positions = sequence.shape
        
        if num_positions != 3:
            raise ValueError(f"序列位置数必须为3，当前为{num_positions}")
        
        # 初始化扩展序列
        extended_sequence = []
        
        # 为每个时间步添加区域转换特征
        for t in range(window_size):
            extended_features = []
            
            # 添加原始数字特征
            for pos in range(3):
                extended_features.append(sequence[t, pos])
            
            # 添加区域转换特征（除了第一个时间步）
            if t > 0:
                for pos in range(3):
                    prev_digit = sequence[t-1, pos]
                    curr_digit = sequence[t, pos]
                    
                    # 计算区域转换
                    # 0-4为小数区，5-9为大数区
                    prev_zone = 0 if prev_digit <= 4 else 1  # 0表示小数区，1表示大数区
                    curr_zone = 0 if curr_digit <= 4 else 1
                    
                    # 记录区域转换：-1表示大数区到小数区，0表示无转换，1表示小数区到大数区
                    transition = curr_zone - prev_zone
                    extended_features.append(transition)
            else:
                # 第一个时间步无区域转换信息，用0填充
                extended_features.extend([0, 0, 0])
            
            extended_sequence.append(extended_features)
        
        return np.array(extended_sequence)
    
    def analyze_zone_transitions(self, data):
        """
        分析区域转换模式统计
        
        Args:
            data: 历史数据，形状为 [num_samples, 3]
            
        Returns:
            区域转换统计字典
        """
        transitions_stats = {
            'small_to_large': 0,  # 小数区到大数区
            'large_to_small': 0,  # 大数区到小数区
            'no_change': 0        # 无区域变化
        }
        
        for i in range(1, len(data)):
            for pos in range(3):
                prev_digit = data[i-1][pos]
                curr_digit = data[i][pos]
                
                prev_zone = 0 if prev_digit <= 4 else 1  # 0表示小数区，1表示大数区
                curr_zone = 0 if curr_digit <= 4 else 1
                
                if prev_zone == 0 and curr_zone == 1:
                    transitions_stats['small_to_large'] += 1
                elif prev_zone == 1 and curr_zone == 0:
                    transitions_stats['large_to_small'] += 1
                else:
                    transitions_stats['no_change'] += 1
        
        return transitions_stats
    
    def prepare_prediction_features(self, recent_data):
        """
        准备用于预测的特征（包含区域转换特征）
        
        Args:
            recent_data: 最近的数据，可以是numpy数组或torch张量
            
        Returns:
            包含区域转换特征的预测数据
        """
        # 确保recent_data是torch.Tensor类型
        if not isinstance(recent_data, torch.Tensor):
            recent_data = torch.tensor(recent_data, dtype=torch.float32)
        
        # 如果是2D张量，添加batch维度
        if len(recent_data.shape) == 2:
            recent_data = recent_data.unsqueeze(0)
        
        # 提取区域转换特征
        if recent_data.shape[-1] == 3:  # 只有原始数字特征
            # 添加区域转换特征
            extended_data = self.extract_zone_transitions(recent_data)
            return torch.FloatTensor(extended_data)
        elif recent_data.shape[-1] == 6:  # 已经包含区域转换特征
            return recent_data
        else:
            raise ValueError(f"不支持的数据维度: {recent_data.shape}")


def load_3d_history_data(csv_file_path):
    """
    加载3D历史数据
    
    Args:
        csv_file_path: CSV文件路径
        
    Returns:
        解析后的数据数组，形状为 [num_samples, 3]
    """
    try:
        # 读取CSV文件
        df = pd.read_csv(csv_file_path)
        
        # 解析num_1, num_2, num_3列
        numbers_data = []
        for _, row in df.iterrows():
            try:
                # 提取3个数字并确保在0-9范围内
                parsed_numbers = []
                for i in range(1, 4):  # num_1, num_2, num_3
                    col_name = f'num_{i}'
                    num = int(row[col_name])
                    # 使用模运算确保在0-9范围内
                    num = num % 10
                    parsed_numbers.append(num)
                numbers_data.append(parsed_numbers)
            except (ValueError, KeyError):
                # 如果有任何问题，添加默认值[0, 0, 0]
                numbers_data.append([0, 0, 0])
        
        return np.array(numbers_data)
    except Exception as e:
        print(f"加载历史数据失败: {e}")
        return np.array([])


# 测试代码
if __name__ == "__main__":
    # 创建特征提取器实例
    extractor = Data3DZoneFeatureExtractor()
    
    # 测试数据
    test_data = np.array([
        [1, 2, 3],  # 第1期
        [6, 7, 8],  # 第2期
        [2, 4, 6],  # 第3期
        [5, 3, 7],  # 第4期
    ])
    
    print("原始数据:")
    print(test_data)
    print(f"形状: {test_data.shape}")
    
    # 提取区域转换特征
    extended_data = extractor.extract_zone_transitions(test_data)
    print("\n包含区域转换特征的数据:")
    print(extended_data)
    print(f"形状: {extended_data.shape}")
    
    # 分析区域转换模式
    stats = extractor.analyze_zone_transitions(test_data)
    print("\n区域转换统计:")
    print(stats)